# -*- coding: utf-8 -*-
# ----------------------------
# @Time    : 2022/5/29 4:23 PM
# @Author  : changqingai
# @FileName: 03-paddle_batch.py
# ----------------------------

import numpy as np
import tensorflow as tf


# ************ 例子1 ****************
def count(stop):
    i = 0
    while True:
        yield i
        i += 1
        if i > stop:
            break


count_ds = tf.data.Dataset.from_generator(count, args=[15],
                                          output_types=tf.int32,
                                          output_shapes=())

batch_count_ds = count_ds.shuffle(20).padded_batch(batch_size=10, drop_remainder=False)
for batch in batch_count_ds.repeat(1):
    print("count_batch:", batch)


batch_data = next(iter(batch_count_ds))
print("batch_data:", batch_data)

print("# ************ 例子2 ****************")
# ************ 例子2 ****************
def gen_series(stop):
    i = 0
    while True:
        size = np.random.randint(1, 5)
        yield i, np.random.normal(size=(size,))
        i += 1
        if i > stop:
            break


series_ds = tf.data.Dataset.from_generator(gen_series, args=[5],
                                           output_types=(tf.int32, tf.float32),
                                           output_shapes=((), (None,)))

series_ds = series_ds.shuffle(2).padded_batch(batch_size=2)
for batch in series_ds.repeat(2):
    print("batch:", batch)


